import os
import os.path as osp
import numpy as np
import random
import matplotlib.pyplot as plt
import collections
import torch
import torchvision
from torch.utils import data
from PIL import Image

# We firstly use vaihingen dataset as source domain.
class ShaoyangDataSet(data.Dataset):
    def __init__(self, root, list_path, max_iters=None, crop_size=(321, 321), mean=(128, 128, 128), scale=True, mirror=True, ignore_label=255, reduce_zero_label=True, set='train'):
        self.root = root
        self.list_path = list_path
        self.crop_size = crop_size
        self.scale = scale
        self.ignore_label = ignore_label
        self.reduce_zero_label = reduce_zero_label
        self.mean = mean
        self.is_mirror = mirror
        self.set = set

        self.img_ids = [i_id.strip() for i_id in open(list_path)]
        if not max_iters==None:
            self.img_ids = self.img_ids * int(np.ceil(float(max_iters) / len(self.img_ids)))
        self.files = []

        # for split in ["train", "trainval", "val"]:
        for name in self.img_ids:
            img_file = osp.join(self.root, "img_dir/%s/%s" % (self.set, name))
            label_file = osp.join(self.root, "ann_dir/%s/%s" % (self.set, name))
            self.files.append({
                "img": img_file,
                "label": label_file,
                "name": name
            })

    def __len__(self):
        return len(self.files)


    def __getitem__(self, index):
        datafiles = self.files[index]

        image = Image.open(datafiles["img"])
        label = Image.open(datafiles["label"])
        name = datafiles["name"]

        # resize
        # image = image.resize(self.crop_size, Image.BICUBIC)
        # label = label.resize(self.crop_size, Image.NEAREST)

        if self.reduce_zero_label:
            label = np.asarray(label, np.uint8)
            # avoid using underflow conversion
            label[label == 0] = 255
            label = label - 1
            label[label == 254] = 255

        image = np.asarray(image, np.float32)
        label = np.asarray(label, np.float32)

        size = image.shape
        # image = image[:, :, ::-1]  # change to BGR
        image -= self.mean  # Order: IR -> R -> G
        image = image.transpose((2, 0, 1))      # Change order: 600x600x3

        return image.copy(), label.copy(), np.array(size), name
